#include "AnaDef.h"
#include "XGLUtils.h"
#include "dEdxGausN.h"

Float_t fv0cen = -999;
Float_t fmceta = -999;  
Float_t fmcpt = -999;  
Int_t fmcpdg = -999;
Int_t fmctype = -999;
Int_t fmccharge = -999;
Int_t fkFailEvent= -999;

Bool_t fkTEST = 0;
Bool_t fkPbPb = kTRUE;
Int_t fkCharge = 0;

Bool_t DoCut()
{  
  //should have primary cut!! otherwise all secondaries are recovered!!
  //select primary
  if(fmctype!=1)
    return kFALSE;


  if(fkPbPb){
    if(!dEdxGausN::IsCenOK(fv0cen))
      return kFALSE;
  }

  if(!dEdxGausN::IsEtaOK(fmceta))
    return kFALSE;

  //fkCharge=0, always true
  //fkCharge=1, fcharge<0 false, i.e. >0 true
  //fkChareg=-1, fchareg>0 false, i.e. <0 true
  if(fmccharge*fkCharge<0)
    return kFALSE;


  return kTRUE;
}

void extractMCgenPtHist(const TString fin)
{
  //input
  TChain *ch=XGLUtils::InputFiles(fin, "mctree");

  ch->SetBranchStatus("*",0);

  ch->SetBranchStatus("mcpdg",1);
  ch->SetBranchAddress("mcpdg",     &fmcpdg);

  ch->SetBranchStatus("mctype",1);
  ch->SetBranchAddress("mctype",     &fmctype);

  ch->SetBranchStatus("mceta",1);
  ch->SetBranchAddress("mceta",     &fmceta);

  ch->SetBranchStatus("mcpt",1);
  ch->SetBranchAddress("mcpt",     &fmcpt);

  ch->SetBranchStatus("v0cen",1);
  ch->SetBranchAddress("v0cen",     &fv0cen);

  ch->SetBranchStatus("mccharge",1);
  ch->SetBranchAddress("mccharge",     &fmccharge);

  ch->SetBranchStatus("kFailEvent",1);
  ch->SetBranchAddress("kFailEvent",     &fkFailEvent); 


  TList *ll=new TList;

  TH2D * PshmcpdgVsPt = dEdxGausN::GetHistMCPDGVsLog10Pt("PshmcpdgVsPt");
  ll->Add(PshmcpdgVsPt);

  TH2D * PsVrecVzhmcpdgVsPt = dEdxGausN::GetHistMCPDGVsLog10Pt("PsVrecVzhmcpdgVsPt");
  ll->Add(PsVrecVzhmcpdgVsPt);

  TH2D * hmcpdgEta = new TH2D("mcpdgVsEta","",200,-1,1, 4,0,4);
  ll->Add(hmcpdgEta);

  TH2D * hmcpwmPt[5];
  for(Int_t ii=0; ii<5; ii++){
    hmcpwmPt[ii]= dEdxGausN::GetHistMCPWMVsLog10Pt(ii);
    ll->Add(hmcpwmPt[ii]);
  }

  TH1I * hmccharge = new TH1I("hmccharge", "", 10, -5, 5);
  ll->Add(hmccharge);

  TH2D * hkFailEventVsPt = dEdxGausN::GetHistMCkFailEventVsLog10Pt();
  ll->Add(hkFailEventVsPt);

  TH1D * hv0cen = new TH1D("hv0cen", "", 200, 0, 100);
  ll->Add(hv0cen);
  
  //original fill!
  Long64_t nent = 0;
  while(ch->GetEntry(nent)){
    if(nent%1000000==0){
      printf("original fill %lld\n", nent);
    }
    
    nent++;

    if(fkTEST){
      if(nent>=1000000)  break;
    }

    if(DoCut()){
      hv0cen->Fill(fv0cen);
      hmccharge->Fill(fmccharge);

      const Int_t itype = dEdxGausN::PDG2Type(fmcpdg);

      PshmcpdgVsPt->Fill(TMath::Log10(fmcpt), itype);

      if(fkFailEvent==0){
        PsVrecVzhmcpdgVsPt->Fill(TMath::Log10(fmcpt), itype);
      }

      hmcpdgEta->Fill(fmceta, itype);

      hkFailEventVsPt->Fill(TMath::Log10(fmcpt), fkFailEvent);

      hmcpwmPt[itype]->Fill(TMath::Log10(fmcpt), fmctype);
      hmcpwmPt[4]->Fill(TMath::Log10(fmcpt), fmctype);
    }
  }

  printf("original fill nent: %lld\n", nent);
  ll->Print();
  ll->ls();
  
  TString tag(fin);
  tag=  tag(tag.Last('/')+1,tag.Last('.')-tag.Last('/')-1);
  tag+="-";

  TString pwd = gSystem->pwd();
  pwd.ReplaceAll("/hera/alice/xlu/train/","");
  pwd.ReplaceAll("output","");
  pwd.ReplaceAll("analysis_To_hist","");
  pwd.ReplaceAll("hadron","");
  pwd.ReplaceAll("/","");

  tag+=pwd;  

  TFile *fout=new TFile(Form("histGen-%s.root", tag.Data()), "recreate");

  ll->Write();
  fout->Save();
  fout->Close();
  
  printf("done!!");
}

int main (int argc, char * argv[])
{
  for(int ii=0; ii<argc; ii++){
    printf("%d: %s\n", ii, argv[ii]);
  }

  if(argc!=2){
    printf("argc!=2\n"); return 1;
  }

  const TString pwd=gSystem->pwd();
  if(!pwd.Contains("MC") || !pwd.Contains("PRIM")){
    printf("not MC nor PRIM! for efficiency PRIM is needed %s\n", pwd.Data());
    return 1;
  }

  dEdxGausN::EtaIni();
  dEdxGausN::CenIni();
 
  if(pwd.Contains("test")) fkTEST = kTRUE;
  else fkTEST = kFALSE;

  if(pwd.Contains("10h") || pwd.Contains("11h")) fkPbPb = kTRUE;
  else fkPbPb = kFALSE;  

  fkCharge = dEdxGausN::GetkCharge();


  printf("fkPbPb %d; fkTEST %d; fkCharge %d\n", fkPbPb, fkTEST, fkCharge);

  dEdxGausN::HistPrint();

  TStopwatch stwa;
  stwa.Start(kTRUE);

  extractMCgenPtHist(argv[1]);

  stwa.Stop(); 
  stwa.Print("m"); //"m": milli sec; "u": micro sec
  
  return 0;
}
